package com.berchen.数据结构.二叉树.平衡二叉树;

import lombok.Data;

/**
 * 平衡二叉树（AVL）
 */
@Data
public class AVLTree {
    private AVLNode root;


    /**
     * 右旋转
     */
    private void rightRotate(AVLNode node) {

        // 1. 以根节点的值创建一个新节点
        AVLNode newNode = new AVLNode(node.value);
        // 2. 将根节点的右子节点赋值给新节点的右子节点
        newNode.right = node.right;
        // 3. 将根节点的左子节点的右子节点赋值给新节点的左子节点
        newNode.left = node.left.right;
        // 4. 将根节点的左子节点的值赋值给跟节点
        node.value = node.left.value;
        // 5. 将新节点赋值给根节点的right
        node.right = newNode;
        // 6. 将根节点的左子节的左子节点赋值给根节点的左子节点
        node.left = node.left.left;
    }

    /**
     * 左旋转
     */
    private void leftRotate(AVLNode node) {

        // 1.以根节点的值创建新的节点
        AVLNode newNode = new AVLNode(node.value);

        // 2.将根节点的左子节点赋值给新节点的left
        newNode.left = node.left;
        // 3.把根节点的右子节点的左子节点赋值给新节点的right
        newNode.right = node.right.left;
        // 4.把根节点的值换成根节点的右子节点的值
        node.value = node.right.value;
        // 5.把根节点的右子节点的右子节点赋值给跟节点的right
        node.right = node.right.right;
        // 6.把新节点赋值给根节点的left
        node.left = newNode;
    }

    /**
     * 获取树的高度
     *
     * @return
     */
    public int getTreeHeight() {
        return root == null ? 0 : Math.max(getHeight(root.left), getHeight(root.right)) + 1;
    }

    /**
     * 获取节点的左子树  或右子树的高度
     *
     * @param flag true 左子树高度  false 右子树高度
     * @param node
     * @return
     */
    public int getHeight(boolean flag, AVLNode node) {
        return node == null ? 0 : flag ? getHeight(node.left) : getHeight(node.right);
    }

    private int getHeight(AVLNode node) {
        return node == null ? 0 : Math.max(node.left == null ? 0 : getHeight(node.left), node.right == null ? 0 : getHeight(node.right)) + 1;
    }

    /**
     * 往二叉排序树中添加元素
     *
     * @param node
     */
    public void add(AVLNode node) {

        if (root == null) {
            root = node;
        } else {
            AVLNode temp = root;
            while (true) {
                if (temp == null) {
                    temp = node;
                } else {
                    // 放右边
                    if (temp.value <= node.value) {
                        if (temp.right == null) {
                            temp.right = node;
                            break;
                        } else {
                            temp = temp.right;
                        }
                    } else {
                        // 放左边
                        if (temp.left == null) {
                            temp.left = node;
                            break;
                        } else {
                            temp = temp.left;
                        }
                    }
                }
            }
        }
        // 当添加完一个节点后 如果右子树高度比左子树高度大于1 就需要左旋转
        if (getHeight(false, root) - getHeight(true, root) > 1) {
            // 如果需要左旋转 的节点的右子树的 左子树 高度 大于 右子树的高度  就先将需要左旋转的节点的右子树进行右旋转
            if(getHeight(true,root.right)>getHeight(false,root.right)){
                rightRotate(root.right);
            }
            // 再进行左选择
            leftRotate(root);
            return; // 已经平衡了
        }
        // 当添加完一个节点后 如果左子树高度比右子树高度大于1  就需要右旋转
        if (getHeight(true, root) - getHeight(false, root) > 1) {
            // 如果需要右旋转 的节点的 左子树的 右子树 高度 大于 左子树的高度  就先将需要右旋转的节点的左子树进行左旋转
            if(getHeight(false,root.left)>getHeight(true,root.left)){
                leftRotate(root.left);
            }
            rightRotate(root);
        }
    }

    /**
     * 删除节点
     *
     * @param value
     */
    public void delNode(int value) {
        if (root == null) {
            return;
        }
        AVLNode targetNode = getDelNode(value);
        if (targetNode == null) {
            return;
        }
        if (root.left == null && root.right == null) {
            root = null;
            return;
        }
        AVLNode delParent = getDelParent(value);
        if (targetNode.left == null && targetNode.right == null) {
            // 删除叶子节点
            if (delParent.left != null && delParent.left.value == value) {
                delParent.left = null;
            }
            if (delParent.right != null && delParent.right.value == value) {
                delParent.right = null;
            }
        } else if (targetNode.left != null && targetNode.right != null) {
            // 删除两个子节点的树
            if (delParent != null) {
                AVLNode minNode = getMinNode(targetNode, false);
                minNode.left = targetNode.left;
                minNode.right = targetNode.right;
                if (delParent.left != null && delParent.left.value == value) {
                    delParent.left = minNode;
                }
                if (delParent.right != null && delParent.right.value == value) {
                    delParent.right = minNode;
                }
            } else {// 父节点为null
                AVLNode minNode = getMinNode(targetNode, false);
                minNode.left = targetNode.left;
                minNode.right = targetNode.right;
                root = minNode;
            }
        } else {
            // 删除一个节点的树
            if (delParent != null) {
                if (delParent.left != null && delParent.left.value == value) {
                    if (targetNode.left != null) {
                        delParent.left = targetNode.left;
                    } else {
                        delParent.left = targetNode.right;
                    }
                } else if (delParent.right != null && delParent.right.value == value) {
                    if (targetNode.left != null) {
                        delParent.right = targetNode.left;
                    } else {
                        delParent.right = targetNode.right;
                    }
                }
            } else {// 父节点为null

                if (targetNode.left != null && targetNode.left.value == value) {
                    root = targetNode.left;
                }
                if (targetNode.right != null && targetNode.right.value == value) {
                    root = targetNode.right;
                }
            }
        }
    }

    /**
     * 找到node为根节点的树的左子树的最大值 还是找到node为根节点的树的右子树的最小值 并且同时删除该节点
     *
     * @param node
     * @param flag true  左子树的最大值   false 右子树的最小值
     * @return
     */
    public AVLNode getMinNode(AVLNode node, boolean flag) {

        if (node == null) {
            return null;
        }
        AVLNode temp = node;
        if (flag) {
            if (temp.left != null) {
                temp = temp.left;
                while (temp.right != null) {
                    temp = temp.right;
                }
                delNode(temp.value);
                return temp;
            }
        } else {
            if (temp.right != null) {
                temp = temp.right;
                while (temp.left != null) {
                    temp = temp.left;
                }
                delNode(temp.value);
                return temp;
            }
        }
        return null;
    }

    /**
     * 获取要删除节点的父节点
     *
     * @param value
     * @return
     */
    public AVLNode getDelParent(int value) {
        if (root == null || root.value == value) {
            return null;
        }
        return getDelParent(root, value);
    }

    private AVLNode getDelParent(AVLNode node, int value) {
        if (node == null) {
            return null;
        }
        if ((node.left != null && node.left.value == value) || (node.right != null && node.right.value == value)) {
            return node;
        }
        AVLNode temp = null;
        if (node.left != null) {
            temp = getDelParent(node.left, value);
        }
        if (temp == null && node.right != null) {
            temp = getDelParent(node.right, value);
        }
        return temp;
    }

    /**
     * 获取要删除的节点
     *
     * @param value
     * @return
     */
    public AVLNode getDelNode(int value) {

        return getDelNode(root, value);
    }

    private AVLNode getDelNode(AVLNode node, int value) {
        if (node == null) {
            return null;
        }
        if (node.value == value) {
            return node;
        }
        AVLNode temp = null;
        if (node.left != null) {
            temp = getDelNode(node.left, value);
        }
        if (temp == null && node.right != null) {
            temp = getDelNode(node.right, value);
        }
        return temp;
    }

    /**
     * 中序遍历
     */
    public void infixOrder() {
        if (root == null) {
            System.out.println("二叉排序树为null");
        } else {
            infixOrder(root);
        }
    }

    private void infixOrder(AVLNode node) {

        if (node.left != null) {
            infixOrder(node.left);
        }
        System.out.println(node);
        if (node.right != null) {
            infixOrder(node.right);
        }
    }
}
